﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Net.Sockets;
using System.Net;
using System.Threading;
using System.Collections;
using UdpSocketClientAsync.Enums;

namespace UdpSocketClientAsync {

    public partial class UdpClientAsync {

        private SocketAsyncEventArgsPool SaeaPool;
        private Semaphore SemaphoreConnections;
        private BufferManager BufferManager;
        private Random RandomSrcPort;
        private SocketStats SocketStats;

        #region Properties

        public int SocketPoolCount {
            get {
                return this.SaeaPool.Count;
            }
        }

        private int _maxConnection;
        public int MaxConnection {
            get {
                return this._maxConnection;
            }
        }

        private int _bufferSize;
        internal int BufferSize {
            get {
                return this._bufferSize;
            }
        }

        private static byte[] _portInUse = new byte[65535];
        public IEnumerable<int> LockedPorts {
            get {
                return UdpClientAsync._portInUse
                    .Select( (b, i) => { return new { Port = i, Locked = b}; })
                    .Where( p => p.Locked == (byte)PortInUseEnum.Locked)
                    .Select( p => { return p.Port; } );

            }
        }

        private TimeSpan _timeout;
        internal TimeSpan Timeout {
            get {
                return this._timeout;
            }
        }

        private bool _useRandomSrcPort;
        internal bool UseRandomSrcPort {
            get {
                return this._useRandomSrcPort;
            }
        }

        private bool _connectRemoteHost;
        private bool ConnectRemoteHost {
            get {
                return this._connectRemoteHost;
            }
        }

        #endregion

        #region SocketStats

        public int TotalSendTo {
            get {
                return this.SocketStats.TotalSendTo;
            }
        }

        public int TotalReceiveFrom {
            get {
                return this.SocketStats.TotalReceiveFrom;
            }
        }

        public int TotalReceiveTimeout {
            get {
                return this.SocketStats.TotalReceiveTimeout;
            }
        }

        public int ResponseSuccess {
            get {
                return this.SocketStats.ResponseSuccess;
            }
        }

        public int AdressAlreadyInUse {
            get {
                return this.SocketStats.AdressAlreadyInUse;
            }
        }

        public int PortCollision {
            get {
                return this.SocketStats.PortCollision;
            }
        }

        public DateTime StartedAt {
            get {
                return this.SocketStats.StartedAt;
            }
        }

        #endregion

        const int DefaultMaxConnection = 10000;
        const int DefaultBufferSize = 512;
        static readonly TimeSpan DefaultTimeout = TimeSpan.FromMilliseconds(3000);
        private int MinSrcPort = 1;
        private int MaxSrcPort = 49152;

        #region Event

        public event UdpClientDelegate OnSendSubmitted;
        public event UdpClientDelegate OnSendCompleted;
        public event UdpClientDelegate OnReceiveCompleted;

        private EventHandler<SocketAsyncEventArgs> sendToCompleted;
        private EventHandler<SocketAsyncEventArgs> receiveFromCompleted;


        #endregion

        #region Ctor

        public UdpClientAsync() : this(DefaultMaxConnection, DefaultBufferSize, DefaultTimeout) {
        }

        public UdpClientAsync(int maxConnection, int bufferSize, TimeSpan timeout) {

            this._maxConnection = maxConnection;
            this._bufferSize = bufferSize;
            this._timeout = timeout;
            this._connectRemoteHost = true;
            this._useRandomSrcPort = true;

            this.SaeaPool = new SocketAsyncEventArgsPool(maxConnection);
            this.SemaphoreConnections = new Semaphore(maxConnection, maxConnection);
            this.BufferManager = new BufferManager(maxConnection * bufferSize, bufferSize);
            this.InitSaeaPool(maxConnection);
            this.RandomSrcPort = new Random();
            this.SocketStats = new SocketStats();

            this.sendToCompleted = new EventHandler<SocketAsyncEventArgs>(saea_SendToCompleted);
            this.receiveFromCompleted = new EventHandler<SocketAsyncEventArgs>(saea_ReceiveFromCompleted);
        }

        private void InitSaeaPool(int maxConnection) {

            // Preallocate pool of SocketAsyncEventArgs objects.
            for (int i = 0; i < maxConnection; i++) {
                SocketAsyncEventArgs e = new SocketAsyncEventArgs();
                this.BufferManager.SetBuffer(e);
                this.SaeaPool.Push(e);
            }

        }

        #endregion

        #region SendTo

        public void SendTo(IPEndPoint remote, byte[] buffer, int offset, int length, object ticket) {

            if (remote == null) {
                throw new ArgumentNullException("remote");
            }

            if (buffer == null) {
                throw new ArgumentNullException("buffer");
            }

            if (offset < 0 || offset >= length) {
                throw new ArgumentOutOfRangeException("offset");
            }

            if (length < 0 || length > buffer.Length - offset) {
                throw new ArgumentOutOfRangeException("length");
            }

            this.SemaphoreConnections.WaitOne();
            SocketAsyncEventArgs e = this.SaeaPool.Pop();

            UdpStateObject stateObject = new UdpStateObject(
                e,
                remote,
                buffer,
                offset,
                length,
                ticket) ;

            ThreadPool.QueueUserWorkItem(new WaitCallback(UdpSocketSend), stateObject);

        }

        private void UdpSocketSend(object state) {

            if (state == null) {
                throw new ArgumentNullException("state");
            }
            UdpStateObject stateObject = (UdpStateObject)state;
            object ticket = stateObject.Ticket;
            SocketAsyncEventArgs e = stateObject.Saea;

            try {

                this.UdpSocketSend(e, 
                    stateObject.RemoteEndPoint, 
                    stateObject.Buffer, 
                    stateObject.Offset, 
                    stateObject.Length, 
                    ticket);

                if (this.OnSendSubmitted != null) {
                    UdpClientEventArgs ucea = new UdpClientEventArgs(ticket);
                    this.OnSendSubmitted(this, ucea);
                }

            }
            catch (Exception ex) {
                
                this.ReleaseSaea(e);

                if (this.OnSendSubmitted != null) {
                    UdpClientEventArgs ucea = new UdpClientEventArgs(ticket, ex);
                    this.OnSendSubmitted(this, ucea);
                    if (!ucea.Treated) {
                        throw;
                    }
                }
            }
        }

        /// <summary>
        /// 
        /// </summary>
        /// <param name="e"></param>
        /// <param name="remote"></param>
        /// <param name="buffer"></param>
        /// <param name="offset"></param>
        /// <param name="length"></param>
        /// <param name="ticket"></param>
        private void UdpSocketSend(SocketAsyncEventArgs e, EndPoint remote, byte[] buffer, int offset, int length, object ticket) {

            if (e == null) {
                throw new ArgumentNullException("SocketAsyncEventArgs");
            }

            if (remote == null) {
                throw new ArgumentNullException("EndPoint");
            }

            if (buffer == null) {
                throw new ArgumentNullException("Buffer");
            }

            if (offset < 0 || offset >= length) {
                throw new ArgumentOutOfRangeException("Offset");
            }

            if (length < 0 || length > buffer.Length - offset) {
                throw new ArgumentOutOfRangeException("Length");
            }

            if (length > this.BufferSize) {
                throw new ArgumentOutOfRangeException("Length2");
            }

            this.BufferManager.ResetBuffer(e);
            Array.Copy(buffer, offset, e.Buffer, e.Offset, length);
            e.SetBuffer(e.Offset, buffer.Length);

            e.RemoteEndPoint = remote;
            e.Completed += sendToCompleted;

            Socket socket = new Socket(AddressFamily.InterNetwork, SocketType.Dgram, ProtocolType.Udp);
            if (this.UseRandomSrcPort) {
                this.SetRandomSrcPort(socket);
            }

            if (this.ConnectRemoteHost) {
                socket.Connect(remote);
            }

            UdpUserToken userToken = new UdpUserToken(remote, socket, ticket);
            e.UserToken = userToken;

            bool iopending = socket.SendToAsync(e);
            if (iopending == false) {
                saea_SendToCompleted(this, e);
            }


        }

        private void SetRandomSrcPort(Socket socket) {

            if (socket == null) {
                throw new ArgumentNullException("socket");
            }

            int port = 0;
            int collision = 0;
            int maxtry = 10;
            bool randomPortAssigned = false;

            while (!randomPortAssigned && maxtry >= 0) {
                maxtry--;

                lock (_portInUse) {

                    do {
                        port = this.RandomSrcPort.Next(MinSrcPort, MaxSrcPort);
                        collision++;
                    } while (_portInUse[port] != (byte)PortInUseEnum.Free);

                    if (collision > 1) {
                        this.SocketStats.IncrementPortCollision();
                    }
                    _portInUse[port] = (byte)PortInUseEnum.InUse;
                    EndPoint ep = new IPEndPoint(0, port);

                    try {

                        socket.Bind(ep);
                        randomPortAssigned = true;

                    } catch (SocketException se) {

                        if (se.SocketErrorCode != SocketError.AddressAlreadyInUse) {
                            throw;
                        } else {
                            this.SocketStats.IncrementAdressAlreadyInUse();
                            _portInUse[port] = (byte)PortInUseEnum.Locked;
                        }

                    }
                }
            }

            if (maxtry <= 0) {
                throw new InvalidOperationException("bind");
            }

        }

        #endregion

        #region SendToCompleted

        private void saea_SendToCompleted(object sender, SocketAsyncEventArgs e) {

            if (e == null) {
                throw new ArgumentNullException("SocketAsyncEventArgs");
            }

            UdpUserToken userToken = (UdpUserToken)e.UserToken;
            if (userToken == null) {
                throw new ArgumentNullException("UserToken");
            }

            object ticket = userToken.Ticket;

            try {

                this.SendToCompleted(e);

                if (this.OnSendCompleted != null) {
                    UdpClientEventArgs ucea = new UdpClientEventArgs(ticket);
                    this.OnSendCompleted(this, ucea);
                }

            }
            catch (Exception ex) {
                
                this.ReleaseSaea(e);

                if (this.OnSendCompleted != null) {
                    UdpClientEventArgs ucea = new UdpClientEventArgs(ticket, ex);
                    this.OnSendCompleted(this, ucea);
                    if (!ucea.Treated) {
                        throw;
                    }
                }
            }
        }

        const bool ExecuteOnlyOnce = true;

        private void SendToCompleted(SocketAsyncEventArgs e) {

            if (e == null) {
                throw new ArgumentNullException("SocketAsyncEventArgs");
            }

            if (e.LastOperation != SocketAsyncOperation.SendTo) {
                throw new InvalidOperationException(e.LastOperation.ToString());
            }

            if (e.SocketError != SocketError.Success) {
                throw new SocketException((int)e.SocketError);
            }

            if (e.UserToken == null) {
                throw new ArgumentNullException("UserToken");
            }

            UdpUserToken userToken = (UdpUserToken)e.UserToken;
            if (userToken.Socket == null) {
                throw new ArgumentNullException("Socket");
            }

            Socket socket = userToken.Socket;
            this.SocketStats.IncrementTotalSendTo();

            // do the job
            this.BufferManager.ResetBuffer(e);

            e.Completed -= sendToCompleted;
            e.Completed += receiveFromCompleted;

            bool iopending = socket.ReceiveFromAsync(e);
            if (iopending == false) {
                saea_ReceiveFromCompleted(this, e);
            }

            RegisteredWaitHandle handle = ThreadPool.RegisterWaitForSingleObject(
                userToken.TimeoutEvent,
                new WaitOrTimerCallback(saea_ReceiveFromTimedOut),
                userToken,
                this.Timeout,
                ExecuteOnlyOnce
             );

        }

        #endregion

        #region ReceiveFromTimeout

        private void saea_ReceiveFromTimedOut(object state, bool timedOut) {

            // check da bitch
            if (timedOut == false) {
                return;
            }
            if (state == null) {
                throw new ArgumentNullException("state");
            }
            UdpUserToken userToken = (UdpUserToken)state;
            lock (userToken) {

                if (userToken.SocketClosed) {
                    return;
                }
                if (userToken.Socket == null) {
                    return;
                }

                // update counters, print stats
                this.SocketStats.IncrementTotalReceiveTimeout();

                // close socket will trigger saea_ReceiveFromCompleted with SocketError.OperationAborted
                // seae will be released in saea_ReceiveFromCompleted.
                this.CloseSocket(userToken); 

            }

        }


        #endregion

        #region ReceiveFromCompleted

        private void saea_ReceiveFromCompleted(object sender, SocketAsyncEventArgs e) {

            if (e == null) {
                throw new ArgumentNullException("SocketAsyncEventArgs");
            }

            UdpUserToken userToken = (UdpUserToken)e.UserToken;
            if (userToken == null) {
                throw new ArgumentNullException("UserToken");
            }

            object ticket = userToken.Ticket;

            try {

                this.ReceiveFromCompleted(e);

                if (this.OnReceiveCompleted != null) {
                    UdpClientEventArgs ucea = new UdpClientEventArgs(ticket, e.Buffer, e.Offset, e.BytesTransferred);
                    this.OnReceiveCompleted(this, ucea);
                }


            }
            catch (Exception ex) {


                // userToken.TimoutEvent may be null in cas of na
                // immediate synchronous event.
                // Thus it does not need to timeout.
                if (userToken.TimeoutEvent != null) {
                    userToken.TimeoutEvent.Set();
                }

                if (this.OnReceiveCompleted != null) {
                    UdpClientEventArgs ucea = new UdpClientEventArgs(ticket, ex);
                    this.OnReceiveCompleted(this, ucea);
                    if (!ucea.Treated) {
                        throw;
                    }
                }

            }
            finally {

                this.ReleaseSaea(e);
            }

        }

        private void ReceiveFromCompleted(SocketAsyncEventArgs e) {

            if (e.LastOperation != SocketAsyncOperation.ReceiveFrom) {
                throw new InvalidOperationException(e.LastOperation.ToString());
            }

            if (e.SocketError != SocketError.Success) {
                throw new SocketException((int)e.SocketError);
            }

            if (e.UserToken == null) {
                throw new ArgumentNullException("userToken");
            }

            UdpUserToken userToken = (UdpUserToken)e.UserToken;

            // synchronize
            lock (userToken) {

                if (!userToken.SocketClosed) {

                    if (userToken.Socket == null) {
                        throw new ArgumentNullException("Socket");
                    }

                    // update counters, print stats
                    this.SocketStats.IncrementTotalReceiveFrom();
                    this.SocketStats.IncrementResponseSuccess();

                    // userToken.TimoutEvent may be null in cas of na
                    // immediate synchronous event.
                    // Thus it does not need to timeout.
                    if (userToken.TimeoutEvent != null) {
                        userToken.TimeoutEvent.Set();
                    }

                    this.CloseSocket(userToken);


                }

            }

        }

        #endregion

        #region CloseSocket

        private void CloseSocket(UdpUserToken userToken) {
            

            lock (_portInUse) {

                IPEndPoint ep = (IPEndPoint)userToken.Socket.LocalEndPoint;
                userToken.CloseSocket();

                // Socket bind failed, no localendpoint
                if (ep != null) { 
                    int port = ep.Port;
                    _portInUse[port] = (byte)PortInUseEnum.Free;
                }
            }

        }

        #endregion

        #region ReleaseSaea

        private void ReleaseSaea(SocketAsyncEventArgs e) {

            e.Completed -= receiveFromCompleted;
            e.Completed -= sendToCompleted;

            e.UserToken = null;

            this.BufferManager.ResetBuffer(e);
            this.SaeaPool.Push(e);
            this.SemaphoreConnections.Release();

        }

        #endregion


    }
}
